import cv2
import numpy as np
import pyautogui
import os

# 函数：执行多尺度模板匹配
def multi_scale_template_matching(template, image, scales=[1]):
    found = None

    for scale in scales:
        # 根据尺度调整模板尺寸
        resized_template = cv2.resize(template, None, fx=scale, fy=scale, interpolation=cv2.INTER_CUBIC)
        w, h = resized_template.shape[::-1]

        # 进行模板匹配
        res = cv2.matchTemplate(image, resized_template, cv2.TM_CCOEFF_NORMED)
        _, max_val, _, max_loc = cv2.minMaxLoc(res)

        # 如果找到了一个新的更好匹配，则更新记录的匹配数据
        if found is None or max_val > found[0]:
            found = (max_val, max_loc, scale, w, h)

    return found

# 捕捉屏幕图像
screenshot = pyautogui.screenshot()
screenshot = cv2.cvtColor(np.array(screenshot), cv2.COLOR_RGB2BGR)

# 对屏幕截图进行灰度转换
screen_gray = cv2.cvtColor(screenshot, cv2.COLOR_BGR2GRAY)

# 读取所有扑克牌模板
template_dir = 'templates'
template_files = [f for f in os.listdir(template_dir) if f.endswith('.png')]
templates = [(f, cv2.imread(os.path.join(template_dir, f), 0)) for f in template_files]

# 定义可能的扑克牌尺寸范围（例如，从50%到150%）
scales = np.linspace(1, 2, 3)
print(scales)

# 对每个模板进行多尺度模板匹配
for template_name, template in templates:
    found = multi_scale_template_matching(template, screen_gray, scales)

    if found:
        max_val, max_loc, scale, w, h = found
        if max_val > 0.9:  # 设置阈值
            top_left = max_loc
            bottom_right = (top_left[0] + w, top_left[1] + h)

            # 在原屏幕截图上画出矩形框
            cv2.rectangle(screenshot, top_left, bottom_right, (0, 0, 255), 2)

            # 打印发现的扑克牌信息
            print(f"Detected card {template_name} at position: {top_left}")

# 显示结果
cv2.imshow('Detected Poker Cards', screenshot)
cv2.waitKey(0)
cv2.destroyAllWindows()
